#!/usr/bin/env python3
"""Generate Python bindings for selected XCB protocol elements.

Only the portions of the protocol explicitly requested are generated.
Types referenced by those requests are pulled in transitively.

The emitted code includes the following:

* Enums (as IntEnum classes)
* Typedefs, XID types, and XID unions (as CData subclasses)
* Structs, including replies, but not requests (as ctypes Structs)
  * [Internal] For structs that are variable-length:
    * [Internal] The iterator class
    * [Internal] Initializers for the ctypes "foo_next" function
  * For structs that contain lists:
    * If the list elements are constant length:
      * A "outerstruct_listname" function returning a ctypes array
      * [Internal] Initializers for the "outerstruct_listname" and
        "outerstruct_listname_length" ctypes functions
    * Otherwise (the list elements are variable length):
      * A "outerstruct_listname" function returning a Python list of
        ctypes structs
      * [Internal] Initializers for the "outerstruct_listname_iterator"
        ctypes function
* For requests:
  * [Internal] For those with replies, a cookie class
  * A wrapper function that will block and error-test the request,
    returning the reply struct (or None, for void-returning functions)
"""

from __future__ import annotations

import argparse
import builtins
import keyword
import re
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

from lxml import etree as ET  # noqa: N812 (traditional name)

if TYPE_CHECKING:
    import io
    from collections.abc import Generator, Iterable, Mapping
    from typing import Any, Self

# ----------------------------------------------------------------------------
# Configuration of what we want to generate

TYPES: dict[str, list[str]] = {
    "xproto": ["Setup"],
}

REQUESTS: dict[str, list[str]] = {
    "xproto": [
        "GetGeometry",
        "GetImage",
        "GetProperty",
        # We handle InternAtom specially.
        "NoOperation",
    ],
    "randr": [
        "QueryVersion",
        "GetScreenResources",
        "GetScreenResourcesCurrent",
        "GetCrtcInfo",
    ],
    "render": [
        "QueryVersion",
        "QueryPictFormats",
    ],
    "xfixes": [
        "QueryVersion",
        "GetCursorImage",
    ],
}

# ----------------------------------------------------------------------------
# Constant data used by the generator

PRIMITIVE_CTYPES: dict[str, str] = {
    "CARD8": "c_uint8",
    "CARD16": "c_uint16",
    "CARD32": "c_uint32",
    "CARD64": "c_uint64",
    "INT8": "c_int8",
    "INT16": "c_int16",
    "INT32": "c_int32",
    "INT64": "c_int64",
    "BYTE": "c_uint8",
    "BOOL": "c_uint8",
    "char": "c_char",
    "void": "None",
}

INT_CTYPES = {"c_int8", "c_int16", "c_int32", "c_int64", "c_uint8", "c_uint16", "c_uint32", "c_uint64"}

EIGHT_BIT_TYPES = {
    "c_int8",
    "c_uint8",
    "c_char",
}

XCB_LENGTH_TYPE = "c_int"

RESERVED_NAMES = set(keyword.kwlist) | set(keyword.softkwlist) | set(dir(builtins))

GENERATED_HEADER = """# Auto-generated by gen_xcb_to_py.py - do not edit manually.

# Since many of the generated functions have many parameters, we disable the pylint warning about too many arguments.
# ruff: noqa: PLR0913

from __future__ import annotations

from ctypes import (
    POINTER,
    Array,
    Structure,
    Union,
    _Pointer,
    c_char,
    c_char_p,
    c_double,
    c_float,
    c_int,
    c_int8,
    c_int16,
    c_int32,
    c_int64,
    c_uint,
    c_uint8,
    c_uint16,
    c_uint32,
    c_uint64,
    c_void_p,
)
from enum import IntEnum

from mss.linux.xcbhelpers import (
    LIB,
    XID,
    Connection,
    VoidCookie,
    array_from_xcb,
    initialize_xcb_typed_func,
    list_from_xcb,
)
"""


# ----------------------------------------------------------------------------
# Utility helpers.


class GenerationError(RuntimeError):
    """Raised when the XML describes a construct this generator cannot handle yet."""

    def __init__(self, message: str, *, element: ET._Element | None = None) -> None:
        super().__init__(message)
        self._element = element

    def __str__(self) -> str:
        base = super().__str__()
        if self._element is None:
            return base

        element_base = getattr(self._element, "base", None)
        element_line = getattr(self._element, "sourceline", None)
        if element_base and element_line:
            return f"{element_base}:{element_line}: {base}"
        return base


@contextmanager
def parsing_note(description: str, element: ET._Element | None = None) -> Generator[None]:
    """Context manager to add parsing context to exceptions.

    Use when parsing XML elements to provide better error messages.

    Example:
        with parsing_note("while parsing struct Foo", element):
            ...
    """
    try:
        yield
    except Exception as exc:
        note = description
        if element is not None:
            base = getattr(element, "base", None)
            line = getattr(element, "sourceline", None)
            if base and line:
                note = f"{description} at {base}:{line}"
        exc.add_note(note)
        raise


class LazyDefn:
    """Base class for lazily parsed protocol definitions.

    We lazily parse certain definitions so that we only need to support parsing
    the features that are actually used by the requested types and requests,
    rather than the entire XCB spec.
    """

    def __init__(self, *, protocol: str, name: str, element: ET._Element) -> None:
        self.protocol = protocol
        self.name = name
        self._element = element
        self._parsed = False
        self._parsing = False

    def _parse(self) -> None:
        raise NotImplementedError

    def _ensure_parsed(self) -> None:
        if self._parsed:
            return
        if self._parsing:
            msg = f"Re-entrant parse detected for {self!r}"
            raise RuntimeError(msg)
        self._parsing = True
        try:
            self._parse()
        finally:
            self._parsing = False
        self._parsed = True

    def parse(self) -> None:
        self._ensure_parsed()

    def __getattr__(self, name: str) -> Any:
        self._ensure_parsed()
        if name in self.__dict__:
            return self.__dict__[name]
        msg = f"{type(self).__name__!r} object has no attribute {name!r}"
        raise AttributeError(msg)

    def __repr__(self) -> str:
        return f"{type(self).__name__}({self.protocol!r}, {self.name!r})"


def resolve_primitive(name: str) -> str | None:
    upper = name.upper()
    if upper in PRIMITIVE_CTYPES:
        return PRIMITIVE_CTYPES[upper]
    return PRIMITIVE_CTYPES.get(name)


# ----------------------------------------------------------------------------
# Parsed protocol structures
#
# These are the structures that represent the parsed XML protocol definitions.


@dataclass
class EnumerationItem:
    name: str
    value: int


@dataclass
class EnumDefn:
    protocol: str
    name: str
    items: list[EnumerationItem]


@dataclass
class XidTypeDefn:
    protocol: str
    name: str


@dataclass
class XidUnionDefn(LazyDefn):
    protocol: str
    name: str
    types: list[str]


@dataclass
class TypedefDefn:
    protocol: str
    name: str
    oldname: str


@dataclass
class Field:
    name: str
    type: str
    enum: str | None = None
    mask: str | None = None


@dataclass
class Pad:
    bytes: int | None = None
    align: int | None = None


@dataclass
class ListField:
    name: str
    type: str
    enum: str | None = None


StructMember = Field | Pad | ListField


class StructLikeDefn(LazyDefn):
    """Base class for struct-like definitions.

    This includes structs, requests, and replies, which are all similarly
    structured.
    """

    def __init__(self, protocol: str, name: str, element: ET._Element) -> None:
        super().__init__(protocol=protocol, name=name, element=element)
        # Fields, padding, lists, in their original order
        self.members: list[StructMember]

    @property
    def fields(self) -> list[Field]:
        self._ensure_parsed()
        return [x for x in self.members if isinstance(x, Field)]

    @property
    def lists(self) -> list[ListField]:
        self._ensure_parsed()
        return [x for x in self.members if isinstance(x, ListField)]

    def _parse_child(self, child: ET._Element) -> None:
        """Parse a single child element of the struct-like definition.

        Subclasses are expected to override this to handle additional child
        elements, but should call super() to handle the common ones.
        """
        if isinstance(child, ET._Comment):  # noqa: SLF001
            return
        match child.tag:
            case "field":
                self.members.append(
                    Field(
                        name=child.attrib["name"],
                        type=child.attrib["type"],
                        enum=child.attrib.get("enum"),
                        mask=child.attrib.get("mask"),
                    )
                )
            case "pad":
                self.members.append(parse_pad(child))
            case "list":
                self.members.append(parse_list(child))
            case "doc":
                return
            case _:
                msg = f"Unsupported member {child.tag} in {self.protocol}:{self.name}"
                raise GenerationError(
                    msg,
                    element=child,
                )

    def _parse(self) -> None:
        with parsing_note(f"while parsing {self.protocol}:{self.name}", self._element):
            self.members = []
            for child in self._element:
                self._parse_child(child)


class StructDefn(StructLikeDefn):
    pass


class ReplyDefn(StructLikeDefn):
    # Note that replies don't have their own name, so we use the request name.
    pass


class RequestDefn(StructLikeDefn):
    def __init__(self, protocol: str, name: str, element: ET._Element) -> None:
        super().__init__(protocol=protocol, name=name, element=element)
        self.reply: ReplyDefn | None

    def _parse_child(self, child: ET._Element) -> None:
        if child.tag == "reply":
            self.reply = ReplyDefn(self.protocol, self.name, child)
        else:
            super()._parse_child(child)

    def _parse(self) -> None:
        self.reply = None
        super()._parse()


TypeDefn = XidTypeDefn | XidUnionDefn | TypedefDefn | StructLikeDefn


# ----------------------------------------------------------------------------
# Protocol container and lookups


@dataclass
class ProtocolModule:
    name: str
    version: tuple[int, int] | None
    enums: dict[str, EnumDefn] = field(default_factory=dict)
    types: dict[str, TypeDefn] = field(default_factory=dict)
    requests: dict[str, RequestDefn] = field(default_factory=dict)
    imports: list[str] = field(default_factory=list)


def parse_enum(protocol: str, elem: ET.Element) -> EnumDefn:
    name = elem.attrib["name"]
    with parsing_note(f"while parsing enum {name}", elem):
        items: list[EnumerationItem] = []
        for item in elem.findall("item"):
            if (value := item.find("value")) is not None:
                items.append(EnumerationItem(item.attrib["name"], int(value.text, 0)))
            elif (bit := item.find("bit")) is not None:
                items.append(EnumerationItem(item.attrib["name"], 1 << int(bit.text, 0)))
            else:
                msg = f"Unsupported enum item in {protocol}:{name}:{item}"
                raise GenerationError(
                    msg,
                    element=item,
                )
    return EnumDefn(protocol, name, items)


def parse_xidunion(protocol: str, elem: ET.Element) -> XidUnionDefn:
    name = elem.attrib["name"]
    with parsing_note(f"while parsing xidunion {name}", elem):
        members: list[str] = []
        for child in elem:
            if isinstance(child, ET._Comment):  # noqa: SLF001
                continue
            if child.tag == "type":
                if child.text is None:
                    msg = "xidunion type entry missing text"
                    raise GenerationError(msg, element=child)
                members.append(child.text.strip())
            elif child.tag == "doc":
                continue
            else:
                msg = f"Unsupported xidunion member {child.tag} in {protocol}:{name}"
                raise GenerationError(
                    msg,
                    element=child,
                )
        if not members:
            msg = f"xidunion {protocol}:{name} must include at least one type"
            raise GenerationError(msg, element=elem)
    return XidUnionDefn(protocol, name, members)


def parse_list(elem: ET.Element) -> ListField:
    return ListField(elem.attrib["name"], elem.attrib["type"], elem.attrib.get("enum"))


def parse_pad(elem: ET.Element) -> Pad:
    with parsing_note("while parsing pad", elem):
        bytes_attr = elem.attrib.get("bytes")
        align_attr = elem.attrib.get("align")
        if (bytes_attr is None) == (align_attr is None):
            msg = "Pad must specify exactly one of 'bytes' or 'align'"
            raise GenerationError(msg, element=elem)
        if bytes_attr is not None:
            return Pad(bytes=int(bytes_attr, 0))
        return Pad(align=int(align_attr, 0))


def parse_protocol(path: Path) -> ProtocolModule:  # noqa: PLR0912, PLR0915
    with parsing_note(f"while parsing protocol {path.name}"):
        tree = ET.parse(path)
        root = tree.getroot()
        protocol = root.attrib["header"]
        if "major-version" in root.attrib:
            version = (int(root.attrib["major-version"]), int(root.attrib["minor-version"]))
        else:
            version = None
        module = ProtocolModule(name=protocol, version=version)
        for child in root:
            if isinstance(child, ET._Comment):  # noqa: SLF001
                continue
            match child.tag:
                case "enum":
                    if child.attrib["name"] in module.enums:
                        msg = f"Duplicate enum {child.attrib['name']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.enums[child.attrib["name"]] = parse_enum(protocol, child)
                case "typedef":
                    if child.attrib["newname"] in module.types:
                        msg = f"Duplicate type {child.attrib['newname']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.types[child.attrib["newname"]] = TypedefDefn(
                        protocol, child.attrib["newname"], child.attrib["oldname"]
                    )
                case "xidtype":
                    if child.attrib["name"] in module.types:
                        msg = f"Duplicate type {child.attrib['name']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.types[child.attrib["name"]] = XidTypeDefn(protocol, child.attrib["name"])
                case "xidunion":
                    if child.attrib["name"] in module.types:
                        msg = f"Duplicate type {child.attrib['name']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.types[child.attrib["name"]] = parse_xidunion(protocol, child)
                case "struct":
                    if child.attrib["name"] in module.types:
                        msg = f"Duplicate type {child.attrib['name']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.types[child.attrib["name"]] = StructDefn(protocol, child.attrib["name"], child)
                case "request":
                    if child.attrib["name"] in module.requests:
                        msg = f"Duplicate request {child.attrib['name']} in protocol {protocol}"
                        raise GenerationError(
                            msg,
                            element=child,
                        )
                    module.requests[child.attrib["name"]] = RequestDefn(protocol, child.attrib["name"], child)
                case "import":
                    # There's actually some leeway in how the imports are resolved.  We only require the imported
                    # module to have been loaded if we need to check it for a type or enum.  Since nearly everything
                    # is loaded from the same file or from xproto, it's rarely needed to do that explicitly.
                    module.imports.append(child.text.strip())
                case "union":
                    # We presently just don't use any unions (just xidunion). If they get used by something else,
                    # we'll end up raising an error at that time.
                    pass
                case "error" | "errorcopy":
                    # We don't need any specialized error data.
                    pass
                case "event" | "eventcopy":
                    # We don't use any events at present.
                    pass
                case _:
                    msg = f"Unknown element {child.tag} in protocol {protocol}"
                    raise GenerationError(msg, element=child)
        return module


class ProtocolRegistry:
    """Holds every protocol module and provides lookup helpers."""

    # This gets passed around a lot.  It might be better to put it in a contextvar, if it gets burdensome.

    def __init__(self, proto_dir: Path) -> None:
        self.modules: dict[str, ProtocolModule] = {}
        self._load_all(proto_dir)

    def _load_all(self, proto_dir: Path) -> None:
        for path in sorted(proto_dir.glob("*.xml")):
            module = parse_protocol(path)
            self.modules[module.name] = module

    def resolve_type(self, protocol: str, name: str) -> TypeDefn:
        # Prefer the supplied protocol, then imports.
        module = self.modules.get(protocol)
        if not module:
            msg = f"Unknown protocol {protocol} when resolving type {name}"
            raise GenerationError(msg)
        if name in module.types:
            return module.types[name]
        for imported_modname in module.imports:
            imported_module = self.modules.get(imported_modname)
            if imported_module is None:
                msg = f"Module {protocol} imports {imported_modname}, which is not loaded"
                raise GenerationError(msg)
            if name in imported_module.types:
                return imported_module.types[name]
        msg = f"Unknown type {name} referenced from {protocol}"
        raise GenerationError(msg)

    def resolve_enum(self, protocol: str, name: str) -> EnumDefn:
        module = self.modules.get(protocol)
        if not module:
            msg = f"Unknown protocol {protocol} when resolving enum {name}"
            raise GenerationError(msg)
        if name in module.enums:
            return module.enums[name]
        for imported_modname in module.imports:
            imported_module = self.modules.get(imported_modname)
            if imported_module is None:
                msg = f"Module {protocol} imports {imported_modname}, which is not loaded"
                raise GenerationError(msg)
            if name in imported_module.enums:
                return imported_module.enums[name]
        msg = f"Unknown enum {name} referenced from {protocol}"
        raise GenerationError(msg)

    def resolve_request(self, protocol: str, name: str) -> RequestDefn:
        if protocol not in self.modules:
            msg = f"Unknown protocol {protocol} when resolving request {name}"
            raise GenerationError(msg)
        rv = self.modules[protocol].requests.get(name)
        if rv is None:
            msg = f"Request {protocol}:{name} not found"
            raise GenerationError(msg)
        return rv


# ----------------------------------------------------------------------------
# Dependency analysis


@dataclass
class TopoSortResult:
    enums: list[EnumDefn]
    types: list[TypeDefn]
    requests: list[RequestDefn]


def toposort_requirements(
    registry: ProtocolRegistry,
    type_requirements: Mapping[str, list[str]],
    request_requirements: Mapping[str, list[str]],
) -> TopoSortResult:
    rv = TopoSortResult([], [], [])
    seen_types: set[tuple[str, str]] = set()

    def appendnew[T](collection: list[T], item: T) -> None:
        if item not in collection:
            collection.append(item)

    def require_member(protocol: str, member: StructMember) -> None:
        if isinstance(member, (Field, ListField)):
            require_type(protocol, member.type)
            if member.enum:
                enum = registry.resolve_enum(protocol, member.enum)
                appendnew(rv.enums, enum)
        elif isinstance(member, Pad):
            pass
        else:
            msg = f"Unrecognized struct member {member}"
            raise GenerationError(msg)

    def require_structlike(protocol: str, entry: StructLikeDefn) -> None:
        for member in entry.members:
            require_member(protocol, member)

    def require_type(protocol: str, name: str) -> None:
        primitive = resolve_primitive(name)
        if primitive:
            return
        entry = registry.resolve_type(protocol, name)
        require_resolved_type(entry)

    def require_resolved_type(entry: TypeDefn) -> None:
        key = (entry.protocol, entry.name)
        if key in seen_types:
            return
        seen_types.add(key)
        if isinstance(entry, XidUnionDefn):
            # We put the union first as an XID, so that the subtypes can be derived from it.
            appendnew(rv.types, entry)
            for typ in entry.types:
                require_type(protocol, typ)
        elif isinstance(entry, StructLikeDefn):
            require_structlike(entry.protocol, entry)
            # The request types should all be handled by a different mechanism.
            assert not isinstance(entry, RequestDefn)  # noqa: S101
            appendnew(rv.types, entry)
        else:
            appendnew(rv.types, entry)

    for protocol, names in type_requirements.items():
        for name in names:
            require_type(protocol, name)

    for protocol, names in request_requirements.items():
        for name in names:
            request = registry.resolve_request(protocol, name)
            require_structlike(protocol, request)
            if request.reply:
                require_resolved_type(request.reply)
            appendnew(rv.requests, request)

    return rv


# ----------------------------------------------------------------------------
# Code generation


@dataclass
class FuncDecl:
    protocol: str
    name: str
    argtypes: list[str]
    restype: str


class CodeWriter:
    def __init__(self, fh: io.TextIOBase) -> None:
        self._fh = fh
        self._indent = 0

    def write(self, line: str = "") -> None:
        if line:
            self._fh.write("    " * self._indent + line + "\n")
        else:
            self._fh.write("\n")

    @contextmanager
    def indent(self) -> Generator[Self]:
        self._indent += 1
        yield self
        self._indent -= 1


# Utilities


def type_is_variable(type_: TypeDefn) -> bool:
    return isinstance(type_, StructLikeDefn) and bool(type_.lists)


def is_eight_bit(registry: ProtocolRegistry, protocol: str, name: str) -> bool:
    primitive = resolve_primitive(name)
    if primitive:
        return primitive in EIGHT_BIT_TYPES
    defn = registry.resolve_type(protocol, name)
    if isinstance(defn, TypedefDefn):
        return is_eight_bit(registry, defn.protocol, defn.oldname)
    return False


def lib_for_proto(protocol: str) -> str:
    if protocol == "xproto":
        return "xcb"
    return protocol


# Naming


def camel_case(name: str, protocol: str | None = None) -> str:
    prefix = "" if protocol in {"xproto", None} else camel_case(protocol)  # type: ignore[arg-type]
    if not name.isupper() and not name.islower():
        # It's already in camel case.
        return prefix + name
    camel_name = name.title().replace("_", "")
    return prefix + camel_name


def snake_case(name: str, protocol: str | None = None) -> str:
    prefix = "" if protocol in {"xproto", None} else f"{snake_case(protocol)}_"  # type: ignore[arg-type]
    if name.islower():
        return prefix + name
    s1 = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", name)
    s2 = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", s1)
    s3 = s2.lower()
    return f"{prefix}{s3}"


def format_enum_name(enum: EnumDefn) -> str:
    """Format an enum class name suitable for Python.

    libxcb doesn't define enum names, but we don't use its .h files
    anyway.

    XCB enums are usually already in camel case.  We usually just
    prepend the extension name if it's not xproto.

    Examples:
    * xproto VisualClass -> VisualClass
    * randr Rotation -> RandrRotation
    """
    return camel_case(enum.name, enum.protocol)


def format_enum_item_name(enum_item: EnumerationItem) -> str:
    """Format an entry name in an enum.

    XCB enums are typically already in camel case, which we preserve.
    If there are already both upper and lower case, then we also
    preserve underscores.

    Examles:
    * DirectColor -> DirectColor
    * Rotate_0 -> Rotate_0
    """
    rv = camel_case(enum_item.name)
    return rv if rv not in RESERVED_NAMES else f"{rv}_"


def format_type_name(typedefn: TypeDefn) -> str:
    """Format a type name suitable for Python.

    libxcb defines type names with the C snake case convention, but we
    don't use its .h files anyway.

    We will change all-caps to title case, and prepend the extension
    if it's not xproto.

    Examples:
    * VISUALTYPE -> Visualtype
    * SetupFailed -> SetupFailed
    * ScreenSize -> RandrScreenSize
    """
    base_name = camel_case(typedefn.name, typedefn.protocol)
    if isinstance(typedefn, ReplyDefn):
        return f"{base_name}Reply"
    return base_name


def format_field_name(field: Field) -> str:
    name = field.name
    return f"{name}_" if name in RESERVED_NAMES else name


def format_function_name(name: str, protocol: str | None = None) -> str:
    return snake_case(name, protocol)


# Version constants


def emit_versions(writer: CodeWriter, registry: ProtocolRegistry) -> None:
    writer.write()
    for module in registry.modules.values():
        if module.version is None:
            continue
        const_prefix = module.name.upper()
        writer.write(f"{const_prefix}_MAJOR_VERSION = {module.version[0]}")
        writer.write(f"{const_prefix}_MINOR_VERSION = {module.version[1]}")


# Enums


def emit_enums(writer: CodeWriter, _registry: ProtocolRegistry, enums: Iterable[EnumDefn]) -> None:
    enums = sorted(enums, key=lambda x: (x.protocol, x.name))
    writer.write()
    writer.write("# Enum classes")
    for defn in enums:
        with parsing_note(f"while emitting enum {defn.protocol}:{defn.name}"):
            class_name = format_enum_name(defn)
            writer.write()
            writer.write(f"class {class_name}(IntEnum):")
            with writer.indent():
                for item in defn.items:
                    item_name = format_enum_item_name(item)
                    writer.write(f"{item_name} = {item.value}")


# Simple (non-struct-like) types


def python_type_for(registry: ProtocolRegistry, protocol: str, name: str) -> str:
    primitive = resolve_primitive(name)
    if primitive:
        return primitive
    entry = registry.resolve_type(protocol, name)
    return format_type_name(entry)


# xidtypes interact with xidunions in kind of a backwards order.  This is to make it possible to pass a Window to a
# function that expects a Drawable.  Example output:
# .   class Drawable(XID): pass
# .   class Window(Drawable): pass
# .   class Pixmap(Drawable): pass
# We can't use "Drawable = Window | Pixmap" because ctypes doesn't know what to do with that when used in argtypes.
def emit_xid(
    writer: CodeWriter,
    _registry: ProtocolRegistry,
    entry: XidTypeDefn | XidUnionDefn,
    derived_from: XidUnionDefn | None,
) -> None:
    class_name = format_type_name(entry)
    derived_from_name = format_type_name(derived_from) if derived_from is not None else "XID"
    writer.write()
    writer.write(f"class {class_name}({derived_from_name}):")
    with writer.indent():
        writer.write("pass")


def emit_xidunion(writer: CodeWriter, registry: ProtocolRegistry, entry: XidUnionDefn) -> None:
    emit_xid(writer, registry, entry, None)


def emit_typedef(writer: CodeWriter, registry: ProtocolRegistry, entry: TypedefDefn) -> None:
    class_name = format_type_name(entry)
    base = python_type_for(registry, entry.protocol, entry.oldname)
    writer.write()
    writer.write(f"class {class_name}({base}):")
    with writer.indent():
        writer.write("pass")


# Struct-like types


def emit_structlike(
    writer: CodeWriter,
    registry: ProtocolRegistry,
    entry: StructLikeDefn,
    members: list[tuple[str, str] | StructMember] | None = None,
) -> list[FuncDecl]:
    class_name = format_type_name(entry)
    rv: list[FuncDecl] = []

    # The member list can be overridden by the caller: a reply structure needs to have the generic reply structure
    # (like the sequence number) alongside it, and the padding byte may cause reordering.
    if members is None:
        members = entry.members  # type: ignore[assignment]
    assert members is not None  # noqa: S101

    writer.write()
    writer.write(f"class {class_name}(Structure):")
    with writer.indent():
        # Fields are name, python_type
        field_entries: list[tuple[str, str]] = []
        seen_list: bool = False
        pad_index = 0
        for member in members:
            if isinstance(member, tuple):
                field_entries.append(member)
            elif isinstance(member, Field):
                if seen_list:
                    msg = f"Structure {entry.protocol}:{entry.name} has fields after lists, which is unsupported"
                    raise GenerationError(
                        msg,
                    )
                name = format_field_name(member)
                type_expr = python_type_for(registry, entry.protocol, member.type)
                field_entries.append((name, type_expr))
            elif isinstance(member, Pad):
                if seen_list:
                    continue
                if member.align is not None or member.bytes is None:
                    msg = f"Struct {entry.protocol}:{entry.name} uses align-based padding, which is unsupported"
                    raise GenerationError(
                        msg,
                    )
                name = f"pad{pad_index}"
                pad_index += 1
                field_entries.append((name, f"c_uint8 * {member.bytes}"))
            elif isinstance(member, ListField):
                # At this stage, we don't need to prepare the libxcb list accessor initializers.  We'll do that when
                # we emit the Python wrappers.
                seen_list = True
            else:
                msg = f"Struct {entry.protocol}:{entry.name} has unrecognized member {member}"
                raise GenerationError(msg)
        assert bool(entry.lists) == seen_list  # noqa: S101

        writer.write("_fields_ = (")
        with writer.indent():
            for name, type_expr in field_entries:
                writer.write(f'("{name}", {type_expr}),')
        writer.write(")")

    if seen_list and not isinstance(entry, ReplyDefn):
        # This is a variable-length structure, and it's presumably being accessed from a containing structure's list.
        # It'll need an iterator type in Python, and has a foo_next function in libxcb.
        #
        # We don't try to determine if it's actually being accessed by a containing structure's list: the only way we
        # get here without a parent structure is if this is in TYPES.  But libxcb still defines the iterator and
        # xcb_setup_next function, so we don't have to track if that's happened.
        iterator_name = f"{class_name}Iterator"
        writer.write()
        writer.write(f"class {iterator_name}(Structure):")
        with writer.indent():
            writer.write('_fields_ = (("data", POINTER(' + class_name + ')), ("rem", c_int), ("index", c_int))')
        next_func_name = f"xcb_{format_function_name(entry.name, entry.protocol)}_next"
        rv.append(FuncDecl(entry.protocol, next_func_name, [f"POINTER({iterator_name})"], "None"))

    return rv


def emit_reply(writer: CodeWriter, registry: ProtocolRegistry, entry: ReplyDefn) -> list[FuncDecl]:
    # Replies have a generic structure at the beginning that isn't in the XML spec:
    #     uint8_t        response_type;
    #     uint8_t        pad0;
    #     uint16_t       sequence;
    #     uint32_t       length;
    # However, if the first field of the reply contents is a single byte, then it replaces pad0 in that structure.
    members = entry.members[:]
    field_entries: list[tuple[str, str] | StructMember] = [("response_type", "c_uint8")]
    if members and isinstance(members[0], Field) and is_eight_bit(registry, entry.protocol, members[0].type):
        member = members.pop(0)
        assert isinstance(member, Field)  # noqa: S101
        name = format_field_name(member)
        type_expr = python_type_for(registry, entry.protocol, member.type)
        field_entries.append((name, type_expr))
    elif members and (isinstance(members[0], Pad) and members[0].bytes == 1):
        # XFixes puts the padding byte explicitly at the start of the replies, but it just gets folded in the same way.
        member = members.pop(0)
        field_entries.append(member)
    else:
        field_entries.append(Pad(bytes=1))
    field_entries.append(("sequence", "c_uint16"))
    field_entries.append(("length", "c_uint32"))
    field_entries += members
    return emit_structlike(writer, registry, entry, field_entries)


# Types


def emit_types(
    writer: CodeWriter,
    registry: ProtocolRegistry,
    types: list[TypeDefn],
) -> list[FuncDecl]:
    rv: list[FuncDecl] = []
    writer.write()
    writer.write("# Generated ctypes structures")

    xid_derived_from: dict[tuple[str, str], XidUnionDefn] = {}
    # We have to emit the unions first, so that we can emit the types that can comprise them as subtypes.
    for union in types:
        if not isinstance(union, XidUnionDefn):
            continue
        emit_xidunion(writer, registry, union)
        for subtype_name in union.types:
            subtype = registry.resolve_type(union.protocol, subtype_name)
            subtype_key = (subtype.protocol, subtype.name)
            if subtype_key in xid_derived_from:
                # We could probably use multiple inheritance, but I don't have a test case.
                msg = (
                    f"XID {subtype.protocol}.{subtype.name} is used in multiple unions.  This is"
                    "not currently supported."
                )
                raise GenerationError(msg)
            xid_derived_from[subtype_key] = union

    for typ in types:
        with parsing_note(f"while emitting type {typ.protocol}:{typ.name}"):
            if isinstance(typ, XidTypeDefn):
                emit_xid(writer, registry, typ, xid_derived_from.get((typ.protocol, typ.name)))
            elif isinstance(typ, XidUnionDefn):
                pass
            elif isinstance(typ, TypedefDefn):
                emit_typedef(writer, registry, typ)
            elif isinstance(typ, StructDefn):
                rv += emit_structlike(writer, registry, typ)
            elif isinstance(typ, ReplyDefn):
                rv += emit_reply(writer, registry, typ)
            else:
                msg = f"Unsupported type kind {type(typ).__name__} for {typ.protocol}:{typ.name}"
                raise GenerationError(
                    msg,
                )
    return rv


# List wrappers


def emit_list_field(
    writer: CodeWriter, registry: ProtocolRegistry, struct: StructLikeDefn, field: ListField
) -> list[FuncDecl]:
    protocol = struct.protocol
    base_func_name = f"{format_function_name(struct.name, protocol)}_{format_function_name(field.name)}"
    outer_type_name = format_type_name(struct)
    if field.type == "void":
        # This means that we're getting a void*.  Use an Array[c_char] instead of a c_void_p, so we have the length
        # information.
        inner_is_variable = False
        inner_type_name = "c_char"
    elif field.type in PRIMITIVE_CTYPES:
        inner_is_variable = False
        inner_type_name = PRIMITIVE_CTYPES[field.type]
    else:
        inner_type = registry.resolve_type(protocol, field.type)
        inner_is_variable = type_is_variable(inner_type)
        inner_type_name = format_type_name(inner_type)
    lib = lib_for_proto(protocol)
    if inner_is_variable:
        iterator_type_name = f"{inner_type_name}Iterator"
        xcb_iterator_func_name = f"xcb_{base_func_name}_iterator"
        xcb_next_func_name = f"xcb_{format_function_name(field.type, struct.protocol)}_next"
        writer.write()
        writer.write(f"def {base_func_name}(r: {outer_type_name}) -> list[{inner_type_name}]:")
        with writer.indent():
            writer.write(f"return list_from_xcb(LIB.{lib}.{xcb_iterator_func_name}, LIB.{lib}.{xcb_next_func_name}, r)")
        return [
            FuncDecl(lib, xcb_iterator_func_name, [f"POINTER({outer_type_name})"], iterator_type_name),
            # The "next" function was defined alongside the iterator type.
        ]
    xcb_array_pointer_func_name = f"xcb_{base_func_name}"
    xcb_array_length_func_name = f"xcb_{base_func_name}_length"
    writer.write()
    writer.write(f"def {base_func_name}(r: {outer_type_name}) -> Array[{inner_type_name}]:")
    with writer.indent():
        writer.write(
            f"return array_from_xcb(LIB.{lib}.{xcb_array_pointer_func_name}, LIB.{lib}.{xcb_array_length_func_name}, r)"
        )
    return [
        FuncDecl(lib, xcb_array_pointer_func_name, [f"POINTER({outer_type_name})"], f"POINTER({inner_type_name})"),
        FuncDecl(lib, xcb_array_length_func_name, [f"POINTER({outer_type_name})"], XCB_LENGTH_TYPE),
    ]


def emit_lists(writer: CodeWriter, registry: ProtocolRegistry, types: list[TypeDefn]) -> list[FuncDecl]:
    rv: list[FuncDecl] = []
    for typ in types:
        if not isinstance(typ, StructLikeDefn):
            continue
        for list_field in typ.lists:
            rv += emit_list_field(writer, registry, typ, list_field)
    return rv


# Request wrapper functions


def emit_requests(writer: CodeWriter, registry: ProtocolRegistry, requests: list[RequestDefn]) -> list[FuncDecl | str]:
    rv: list[FuncDecl | str] = []
    for request in requests:
        lib = lib_for_proto(request.protocol)
        func_name = format_function_name(request.name, request.protocol)
        xcb_func_name = f"xcb_{func_name}"
        if request.lists:
            msg = "Cannot handle requests with lists at present"
            raise GenerationError(msg)
        # Parameters are the inputs you declare in the function's "def" line.  Arguments are the inputs you provide to
        # a function when you call it.
        params: list[tuple[str, str]] = [("c", "Connection")]
        params += [
            (format_field_name(field), python_type_for(registry, request.protocol, field.type))
            for field in request.fields
        ]
        params_types = [p[1] for p in params]
        # Arrange for the wrappers to take Python ints in place of any of the int-based ctypes.
        params_with_alts = [(p[0], f"{p[1]} | int" if p[1] in INT_CTYPES else p[1]) for p in params]
        params_string = ", ".join(f"{p[0]}: {p[1]}" for p in params_with_alts)
        args_string = ", ".join(p[0] for p in params)
        if request.reply is None:
            writer.write()
            writer.write(f"def {func_name}({params_string}) -> None:")
            with writer.indent():
                writer.write(f"return LIB.{lib}.{xcb_func_name}({args_string}).check(c)")
            rv.append(FuncDecl(request.protocol, xcb_func_name, params_types, "VoidCookie"))
        else:
            reply_type = request.reply
            reply_type_name = format_type_name(reply_type)
            writer.write()
            writer.write(f"def {func_name}({params_string}) -> {reply_type_name}:")
            with writer.indent():
                writer.write(f"return LIB.{lib}.{xcb_func_name}({args_string}).reply(c)")
            # We have to use initialize_xcb_typed_func to initialize late, rather than making the cookie class here,
            # because the cookie definition needs to reference the XCB reply function.  We could also do a lazy
            # initialization, but it's probably not worth it.
            xcb_params_types = ["POINTER(Connection)", *params_types[1:]]
            rv.append(
                f'initialize_xcb_typed_func(LIB.{lib}, "{xcb_func_name}", '
                f"[{', '.join(xcb_params_types)}], {reply_type_name})"
            )
    return rv


# Initializer function


def emit_initialize(writer: CodeWriter, func_decls: list[FuncDecl | str]) -> None:
    writer.write()
    writer.write("def initialize() -> None:  # noqa: PLR0915")
    with writer.indent():
        for decl in func_decls:
            if isinstance(decl, str):
                writer.write(decl)
            else:
                lib = lib_for_proto(decl.protocol)
                writer.write(f"LIB.{lib}.{decl.name}.argtypes = ({', '.join(decl.argtypes)},)")
                writer.write(f"LIB.{lib}.{decl.name}.restype = {decl.restype}")


# Top level code generator


def generate(
    output: io.TextIOBase,
    proto_dir: Path,
    type_requirements: Mapping[str, list[str]] | None = None,
    request_requirements: Mapping[str, list[str]] | None = None,
) -> None:
    registry = ProtocolRegistry(proto_dir)
    type_requirements = type_requirements or TYPES
    request_requirements = request_requirements or REQUESTS
    plan = toposort_requirements(registry, type_requirements, request_requirements)

    func_decls: list[FuncDecl | str] = []

    writer = CodeWriter(output)
    writer.write(GENERATED_HEADER.rstrip())
    emit_versions(writer, registry)
    emit_enums(writer, registry, plan.enums)
    func_decls += emit_types(writer, registry, plan.types)
    func_decls += emit_lists(writer, registry, plan.types)
    func_decls += emit_requests(writer, registry, plan.requests)
    emit_initialize(writer, func_decls)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Generate ctypes bindings from XCB protocol XML")
    parser.add_argument("--proto_dir", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--output_path", type=Path, default=Path("xcbgen.py"))
    args = parser.parse_args(argv)
    with args.output_path.open("w") as fh:
        generate(fh, args.proto_dir, TYPES, REQUESTS)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
